package com.simple.rag;

import org.neo4j.driver.*;

import java.util.ArrayList;
import java.util.List;

import static org.neo4j.driver.Values.parameters;

/**
 * 项目: middle-ware-design
 * <p>
 * 功能描述:
 *
 * @author: WuChengXing
 * @create: 2025-02-14 00:45
 **/
public class Neo4jRetriever {

    private final Driver driver;

    public Neo4jRetriever(String uri, String user, String password) {
        this.driver = GraphDatabase.driver(uri, AuthTokens.basic(user, password));
    }

    public List<String> retrieveRelatedInfo(String question) {
        List<String> results = new ArrayList<>();
        try (Session session = driver.session()) {
            // 简单示例查询，根据问题关键字查找相关节点属性
            String query = "MATCH (n) WHERE n.name CONTAINS $keyword RETURN n.name";
            Result result = session.run(query, parameters("keyword", question));
            while (result.hasNext()) {
                Record record = result.next();
                results.add(record.get("n.name").asString());
            }
        }
        return results;
    }

    public void close() {
        driver.close();
    }
}
